#!/usr/bin/env python3
import onnx
import onnxruntime
import numpy as np
import os

def optimize_onnx_model(input_path, output_path):
    """优化ONNX模型，尝试消除if节点"""
    
    # 加载原始模型
    print(f"加载ONNX模型: {input_path}")
    model = onnx.load(input_path)
    
    # 分析原始模型
    if_nodes = [node for node in model.graph.node if node.op_type == 'If']
    print(f"原始模型中的If节点数量: {len(if_nodes)}")
    
    # 尝试优化模型
    print("开始优化ONNX模型...")
    
    # 方法1: 使用ONNX的优化器
    try:
        from onnxruntime.transformers import optimizer
        opt_model = optimizer.optimize_model(input_path, model_type='bert')
        opt_model.save_model_to_file(output_path)
        print(f"使用ONNX Runtime优化器优化完成: {output_path}")
    except Exception as e:
        print(f"ONNX Runtime优化器失败: {e}")
        
        # 方法2: 手动移除if节点
        print("尝试手动移除if节点...")
        optimized_model = remove_if_nodes(model)
        onnx.save(optimized_model, output_path)
        print(f"手动优化完成: {output_path}")
    
    # 验证优化后的模型
    if os.path.exists(output_path):
        try:
            optimized_model = onnx.load(output_path)
            if_nodes_after = [node for node in optimized_model.graph.node if node.op_type == 'If']
            print(f"优化后模型中的If节点数量: {len(if_nodes_after)}")
            
            if len(if_nodes_after) == 0:
                print("成功: 所有If节点已被移除!")
            else:
                print("警告: 仍有If节点存在")
                for i, node in enumerate(if_nodes_after):
                    print(f"  If节点 {i}: {node.name}")
                    
        except Exception as e:
            print(f"验证优化后模型失败: {e}")

def remove_if_nodes(model):
    """手动移除if节点，将其替换为更简单的操作"""
    
    # 创建新的图
    new_nodes = []
    
    for node in model.graph.node:
        if node.op_type == 'If':
            print(f"发现If节点: {node.name}")
            
            # 对于if节点，我们尝试用更简单的操作替换
            # 这里我们用一个Identity节点替换，保持输入输出不变
            identity_node = onnx.helper.make_node(
                'Identity',
                inputs=node.input[:1],  # 使用第一个输入
                outputs=node.output,
                name=node.name + '_replaced'
            )
            new_nodes.append(identity_node)
            print(f"  将If节点替换为Identity节点")
        else:
            new_nodes.append(node)
    
    # 创建新的图
    new_graph = onnx.helper.make_graph(
        new_nodes,
        model.graph.name,
        model.graph.input,
        model.graph.output,
        model.graph.initializer
    )
    
    # 创建新的模型
    new_model = onnx.helper.make_model(
        new_graph,
        producer_name=model.producer_name,
        producer_version=model.producer_version,
        ir_version=model.ir_version,
        doc_string=model.doc_string,
        domain=model.domain,
        model_version=model.model_version,
        opset_imports=model.opset_import
    )
    
    return new_model

def main():
    input_path = "/disk2/xd/project/mmdetection/checkpoints/co_detr_custom.onnx"
    output_path = "/disk2/xd/project/mmdetection/checkpoints/co_detr_optimized.onnx"
    
    if not os.path.exists(input_path):
        print(f"输入文件不存在: {input_path}")
        return
    
    optimize_onnx_model(input_path, output_path)
    
    # 检查文件大小
    if os.path.exists(output_path):
        size_mb = os.path.getsize(output_path) / (1024*1024)
        print(f"优化后模型大小: {size_mb:.2f} MB")
    import netron
    netron.start(output_path)

if __name__ == "__main__":
    main() 